import numpy as np
import copy
import open3d as o3d
from PyQt5.QtWidgets import QMessageBox
from uitools import cwd
help_registration = [
   "(1)选择要拼接的两个点云文件，点云以ply/pnt结尾，本操作将把两个不同坐标系下的点云统一起来；" ,
    "(2)程序开始前会让您在同一坐标观测选择的文件，您需要关注两个文件中，那些特征点在世界坐标中是同一个点；",
    "(3)请在点云中选择至少三个共同点，shift+鼠标左击是选择，shift+鼠标右击是取消选择。选中后屏幕会显示一个小球。至少选择3个共同点后，按q退出，选择共同点数量越多越好",
    f"(4)合并后的文件被保存在:{cwd}/process/new_registration.ply"
]




def draw_registration_result(source, target, transformation):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp])


def pick_points(pcd):
    vis = o3d.visualization.VisualizerWithEditing()
    vis.create_window()
    vis.add_geometry(pcd)
    vis.run()  # user picks points
    vis.destroy_window()
    return vis.get_picked_points()


def view_pnt(src,tar):
    source = o3d.io.read_point_cloud(src)
    target = o3d.io.read_point_cloud(tar)
    draw_registration_result(source, target, np.identity(4))

def manual_registration(src,tar):
    source = o3d.io.read_point_cloud(src)
    target = o3d.io.read_point_cloud(tar)

    picked_id_source = pick_points(source)
    picked_id_target = pick_points(target)
    assert (len(picked_id_source) >= 3 and len(picked_id_target) >= 3)
    assert (len(picked_id_source) == len(picked_id_target))
    corr = np.zeros((len(picked_id_source), 2))
    corr[:, 0] = picked_id_source
    corr[:, 1] = picked_id_target

    p2p = o3d.registration.TransformationEstimationPointToPoint()
    trans_init = p2p.compute_transformation(source, target,o3d.utility.Vector2iVector(corr))

    threshold = 0.03  # 3cm distance threshold
    reg_p2p = o3d.registration.registration_icp(
        source, target, threshold, trans_init,
        o3d.registration.TransformationEstimationPointToPoint())
    draw_registration_result(source, target, reg_p2p.transformation)

    new=source.transform(reg_p2p.transformation)+target
    o3d.io.write_point_cloud(cwd+"/process/new_registration.ply", new)


if __name__ == "__main__":
    manual_registration("./raw/box.ply","./raw/office.ply")
